/*
 * Copyright (C) 2007 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.dx.ssa;

import java.util.ArrayList;
import java.util.BitSet;

import com.android.dx.rop.code.CstInsn;
import com.android.dx.rop.code.Insn;
import com.android.dx.rop.code.PlainInsn;
import com.android.dx.rop.code.RegOps;
import com.android.dx.rop.code.RegisterSpec;
import com.android.dx.rop.code.RegisterSpecList;
import com.android.dx.rop.code.Rop;
import com.android.dx.rop.code.Rops;
import com.android.dx.rop.cst.Constant;
import com.android.dx.rop.cst.CstInteger;
import com.android.dx.rop.cst.TypedConstant;
import com.android.dx.rop.type.Type;
import com.android.dx.rop.type.TypeBearer;

/**
 * A small variant of Wegman and Zadeck's Sparse Conditional Constant
 * Propagation algorithm.
 */
public class SCCP {

	/** Lattice values */
	private static final int TOP = 0;
	private static final int CONSTANT = 1;
	private static final int VARYING = 2;
	/** method we're processing */
	private SsaMethod ssaMeth;
	/** ssaMeth.getRegCount() */
	private int regCount;
	/** Lattice values for each SSA register */
	private int[] latticeValues;
	/** For those registers that are constant, this is the constant value */
	private Constant[] latticeConstants;
	/** Worklist of basic blocks to be processed */
	private ArrayList<SsaBasicBlock> cfgWorklist;
	/** Worklist of executed basic blocks with phis to be processed */
	private ArrayList<SsaBasicBlock> cfgPhiWorklist;
	/** Bitset containing bits for each block that has been found executable */
	private BitSet executableBlocks;
	/** Worklist for SSA edges. This is a list of registers to process */
	private ArrayList<SsaInsn> ssaWorklist;
	/**
	 * Worklist for SSA edges that represent varying values. It makes the
	 * algorithm much faster if you move all values to VARYING as fast as
	 * possible.
	 */
	private ArrayList<SsaInsn> varyingWorklist;
	/** Worklist of potential branches to convert to gotos */
	private ArrayList<SsaInsn> branchWorklist;

	private SCCP(SsaMethod ssaMeth) {
		this.ssaMeth = ssaMeth;
		this.regCount = ssaMeth.getRegCount();
		this.latticeValues = new int[this.regCount];
		this.latticeConstants = new Constant[this.regCount];
		this.cfgWorklist = new ArrayList<SsaBasicBlock>();
		this.cfgPhiWorklist = new ArrayList<SsaBasicBlock>();
		this.executableBlocks = new BitSet(ssaMeth.getBlocks().size());
		this.ssaWorklist = new ArrayList<SsaInsn>();
		this.varyingWorklist = new ArrayList<SsaInsn>();
		this.branchWorklist = new ArrayList<SsaInsn>();
		for (int i = 0; i < this.regCount; i++) {
			latticeValues[i] = TOP;
			latticeConstants[i] = null;
		}
	}

	/**
	 * Performs sparse conditional constant propagation on a method.
	 * 
	 * @param ssaMethod
	 *            Method to process
	 */
	public static void process(SsaMethod ssaMethod) {
		new SCCP(ssaMethod).run();
	}

	/**
	 * Adds a SSA basic block to the CFG worklist if it's unexecuted, or to the
	 * CFG phi worklist if it's already executed.
	 * 
	 * @param ssaBlock
	 *            Block to add
	 */
	private void addBlockToWorklist(SsaBasicBlock ssaBlock) {
		if (!executableBlocks.get(ssaBlock.getIndex())) {
			cfgWorklist.add(ssaBlock);
			executableBlocks.set(ssaBlock.getIndex());
		} else {
			cfgPhiWorklist.add(ssaBlock);
		}
	}

	/**
	 * Adds an SSA register's uses to the SSA worklist.
	 * 
	 * @param reg
	 *            SSA register
	 * @param latticeValue
	 *            new lattice value for @param reg.
	 */
	private void addUsersToWorklist(int reg, int latticeValue) {
		if (latticeValue == VARYING) {
			for (SsaInsn insn : ssaMeth.getUseListForRegister(reg)) {
				varyingWorklist.add(insn);
			}
		} else {
			for (SsaInsn insn : ssaMeth.getUseListForRegister(reg)) {
				ssaWorklist.add(insn);
			}
		}
	}

	/**
	 * Sets a lattice value for a register to value.
	 * 
	 * @param reg
	 *            SSA register
	 * @param value
	 *            Lattice value
	 * @param cst
	 *            Constant value (may be null)
	 * @return true if the lattice value changed.
	 */
	private boolean setLatticeValueTo(int reg, int value, Constant cst) {
		if (value != CONSTANT) {
			if (latticeValues[reg] != value) {
				latticeValues[reg] = value;
				return true;
			}
			return false;
		} else {
			if (latticeValues[reg] != value
					|| !latticeConstants[reg].equals(cst)) {
				latticeValues[reg] = value;
				latticeConstants[reg] = cst;
				return true;
			}
			return false;
		}
	}

	/**
	 * Simulates a PHI node and set the lattice for the result to the
	 * appropriate value. Meet values: TOP x anything = TOP VARYING x anything =
	 * VARYING CONSTANT x CONSTANT = CONSTANT if equal constants, VARYING
	 * otherwise
	 * 
	 * @param insn
	 *            PHI to simulate.
	 */
	private void simulatePhi(PhiInsn insn) {
		int phiResultReg = insn.getResult().getReg();

		if (latticeValues[phiResultReg] == VARYING) {
			return;
		}

		RegisterSpecList sources = insn.getSources();
		int phiResultValue = TOP;
		Constant phiConstant = null;
		int sourceSize = sources.size();

		for (int i = 0; i < sourceSize; i++) {
			int predBlockIndex = insn.predBlockIndexForSourcesIndex(i);
			int sourceReg = sources.get(i).getReg();
			int sourceRegValue = latticeValues[sourceReg];

			if (!executableBlocks.get(predBlockIndex)) {
				continue;
			}

			if (sourceRegValue == CONSTANT) {
				if (phiConstant == null) {
					phiConstant = latticeConstants[sourceReg];
					phiResultValue = CONSTANT;
				} else if (!latticeConstants[sourceReg].equals(phiConstant)) {
					phiResultValue = VARYING;
					break;
				}
			} else {
				phiResultValue = sourceRegValue;
				break;
			}
		}
		if (setLatticeValueTo(phiResultReg, phiResultValue, phiConstant)) {
			addUsersToWorklist(phiResultReg, phiResultValue);
		}
	}

	/**
	 * Simulate a block and note the results in the lattice.
	 * 
	 * @param block
	 *            Block to visit
	 */
	private void simulateBlock(SsaBasicBlock block) {
		for (SsaInsn insn : block.getInsns()) {
			if (insn instanceof PhiInsn) {
				simulatePhi((PhiInsn) insn);
			} else {
				simulateStmt(insn);
			}
		}
	}

	/**
	 * Simulate the phis in a block and note the results in the lattice.
	 * 
	 * @param block
	 *            Block to visit
	 */
	private void simulatePhiBlock(SsaBasicBlock block) {
		for (SsaInsn insn : block.getInsns()) {
			if (insn instanceof PhiInsn) {
				simulatePhi((PhiInsn) insn);
			} else {
				return;
			}
		}
	}

	private static String latticeValName(int latticeVal) {
		switch (latticeVal) {
		case TOP:
			return "TOP";
		case CONSTANT:
			return "CONSTANT";
		case VARYING:
			return "VARYING";
		default:
			return "UNKNOWN";
		}
	}

	/**
	 * Simulates branch insns, if possible. Adds reachable successor blocks to
	 * the CFG worklists.
	 * 
	 * @param insn
	 *            branch to simulate
	 */
	private void simulateBranch(SsaInsn insn) {
		Rop opcode = insn.getOpcode();
		RegisterSpecList sources = insn.getSources();

		boolean constantBranch = false;
		boolean constantSuccessor = false;

		// Check if the insn is a branch with a constant condition
		if (opcode.getBranchingness() == Rop.BRANCH_IF) {
			Constant cA = null;
			Constant cB = null;

			RegisterSpec specA = sources.get(0);
			int regA = specA.getReg();
			if (!ssaMeth.isRegALocal(specA) && latticeValues[regA] == CONSTANT) {
				cA = latticeConstants[regA];
			}

			if (sources.size() == 2) {
				RegisterSpec specB = sources.get(1);
				int regB = specB.getReg();
				if (!ssaMeth.isRegALocal(specB)
						&& latticeValues[regB] == CONSTANT) {
					cB = latticeConstants[regB];
				}
			}

			// Calculate the result of the condition
			if (cA != null && sources.size() == 1) {
				switch (((TypedConstant) cA).getBasicType()) {
				case Type.BT_INT:
					constantBranch = true;
					int vA = ((CstInteger) cA).getValue();
					switch (opcode.getOpcode()) {
					case RegOps.IF_EQ:
						constantSuccessor = (vA == 0);
						break;
					case RegOps.IF_NE:
						constantSuccessor = (vA != 0);
						break;
					case RegOps.IF_LT:
						constantSuccessor = (vA < 0);
						break;
					case RegOps.IF_GE:
						constantSuccessor = (vA >= 0);
						break;
					case RegOps.IF_LE:
						constantSuccessor = (vA <= 0);
						break;
					case RegOps.IF_GT:
						constantSuccessor = (vA > 0);
						break;
					default:
						throw new RuntimeException("Unexpected op");
					}
					break;
				default:
					// not yet supported
				}
			} else if (cA != null && cB != null) {
				switch (((TypedConstant) cA).getBasicType()) {
				case Type.BT_INT:
					constantBranch = true;
					int vA = ((CstInteger) cA).getValue();
					int vB = ((CstInteger) cB).getValue();
					switch (opcode.getOpcode()) {
					case RegOps.IF_EQ:
						constantSuccessor = (vA == vB);
						break;
					case RegOps.IF_NE:
						constantSuccessor = (vA != vB);
						break;
					case RegOps.IF_LT:
						constantSuccessor = (vA < vB);
						break;
					case RegOps.IF_GE:
						constantSuccessor = (vA >= vB);
						break;
					case RegOps.IF_LE:
						constantSuccessor = (vA <= vB);
						break;
					case RegOps.IF_GT:
						constantSuccessor = (vA > vB);
						break;
					default:
						throw new RuntimeException("Unexpected op");
					}
					break;
				default:
					// not yet supported
				}
			}
		}

		/*
		 * If condition is constant, add only the target block to the worklist.
		 * Otherwise, add all successors to the worklist.
		 */
		SsaBasicBlock block = insn.getBlock();

		if (constantBranch) {
			int successorBlock;
			if (constantSuccessor) {
				successorBlock = block.getSuccessorList().get(1);
			} else {
				successorBlock = block.getSuccessorList().get(0);
			}
			addBlockToWorklist(ssaMeth.getBlocks().get(successorBlock));
			branchWorklist.add(insn);
		} else {
			for (int i = 0; i < block.getSuccessorList().size(); i++) {
				int successorBlock = block.getSuccessorList().get(i);
				addBlockToWorklist(ssaMeth.getBlocks().get(successorBlock));
			}
		}
	}

	/**
	 * Simulates math insns, if possible.
	 * 
	 * @param insn
	 *            non-null insn to simulate
	 * @param resultType
	 *            basic type of the result
	 * @return constant result or null if not simulatable.
	 */
	private Constant simulateMath(SsaInsn insn, int resultType) {
		Insn ropInsn = insn.getOriginalRopInsn();
		int opcode = insn.getOpcode().getOpcode();
		RegisterSpecList sources = insn.getSources();
		int regA = sources.get(0).getReg();
		Constant cA;
		Constant cB;

		if (latticeValues[regA] != CONSTANT) {
			cA = null;
		} else {
			cA = latticeConstants[regA];
		}

		if (sources.size() == 1) {
			CstInsn cstInsn = (CstInsn) ropInsn;
			cB = cstInsn.getConstant();
		} else { /* sources.size() == 2 */
			int regB = sources.get(1).getReg();
			if (latticeValues[regB] != CONSTANT) {
				cB = null;
			} else {
				cB = latticeConstants[regB];
			}
		}

		if (cA == null || cB == null) {
			// TODO handle a constant of 0 with MUL or AND
			return null;
		}

		switch (resultType) {
		case Type.BT_INT:
			int vR;
			boolean skip = false;

			int vA = ((CstInteger) cA).getValue();
			int vB = ((CstInteger) cB).getValue();

			switch (opcode) {
			case RegOps.ADD:
				vR = vA + vB;
				break;
			case RegOps.SUB:
				// 1 source for reverse sub, 2 sources for regular sub
				if (sources.size() == 1) {
					vR = vB - vA;
				} else {
					vR = vA - vB;
				}
				break;
			case RegOps.MUL:
				vR = vA * vB;
				break;
			case RegOps.DIV:
				if (vB == 0) {
					skip = true;
					vR = 0; // just to hide a warning
				} else {
					vR = vA / vB;
				}
				break;
			case RegOps.AND:
				vR = vA & vB;
				break;
			case RegOps.OR:
				vR = vA | vB;
				break;
			case RegOps.XOR:
				vR = vA ^ vB;
				break;
			case RegOps.SHL:
				vR = vA << vB;
				break;
			case RegOps.SHR:
				vR = vA >> vB;
				break;
			case RegOps.USHR:
				vR = vA >>> vB;
				break;
			case RegOps.REM:
				if (vB == 0) {
					skip = true;
					vR = 0; // just to hide a warning
				} else {
					vR = vA % vB;
				}
				break;
			default:
				throw new RuntimeException("Unexpected op");
			}

			return skip ? null : CstInteger.make(vR);

		default:
			// not yet supported
			return null;
		}
	}

	/**
	 * Simulates a statement and set the result lattice value.
	 * 
	 * @param insn
	 *            instruction to simulate
	 */
	private void simulateStmt(SsaInsn insn) {
		Insn ropInsn = insn.getOriginalRopInsn();
		if (ropInsn.getOpcode().getBranchingness() != Rop.BRANCH_NONE
				|| ropInsn.getOpcode().isCallLike()) {
			simulateBranch(insn);
		}

		int opcode = insn.getOpcode().getOpcode();
		RegisterSpec result = insn.getResult();

		if (result == null) {
			// Find move-result-pseudo result for int div and int rem
			if (opcode == RegOps.DIV || opcode == RegOps.REM) {
				SsaBasicBlock succ = insn.getBlock().getPrimarySuccessor();
				result = succ.getInsns().get(0).getResult();
			} else {
				return;
			}
		}

		int resultReg = result.getReg();
		int resultValue = VARYING;
		Constant resultConstant = null;

		switch (opcode) {
		case RegOps.CONST: {
			CstInsn cstInsn = (CstInsn) ropInsn;
			resultValue = CONSTANT;
			resultConstant = cstInsn.getConstant();
			break;
		}
		case RegOps.MOVE: {
			if (insn.getSources().size() == 1) {
				int sourceReg = insn.getSources().get(0).getReg();
				resultValue = latticeValues[sourceReg];
				resultConstant = latticeConstants[sourceReg];
			}
			break;
		}
		case RegOps.ADD:
		case RegOps.SUB:
		case RegOps.MUL:
		case RegOps.DIV:
		case RegOps.AND:
		case RegOps.OR:
		case RegOps.XOR:
		case RegOps.SHL:
		case RegOps.SHR:
		case RegOps.USHR:
		case RegOps.REM: {
			resultConstant = simulateMath(insn, result.getBasicType());
			if (resultConstant != null) {
				resultValue = CONSTANT;
			}
			break;
		}
		case RegOps.MOVE_RESULT_PSEUDO: {
			if (latticeValues[resultReg] == CONSTANT) {
				resultValue = latticeValues[resultReg];
				resultConstant = latticeConstants[resultReg];
			}
			break;
		}
		// TODO: Handle non-int arithmetic.
		// TODO: Eliminate check casts that we can prove the type of.
		default: {
		}
		}
		if (setLatticeValueTo(resultReg, resultValue, resultConstant)) {
			addUsersToWorklist(resultReg, resultValue);
		}
	}

	private void run() {
		SsaBasicBlock firstBlock = ssaMeth.getEntryBlock();
		addBlockToWorklist(firstBlock);

		/* Empty all the worklists by propagating our values */
		while (!cfgWorklist.isEmpty() || !cfgPhiWorklist.isEmpty()
				|| !ssaWorklist.isEmpty() || !varyingWorklist.isEmpty()) {
			while (!cfgWorklist.isEmpty()) {
				int listSize = cfgWorklist.size() - 1;
				SsaBasicBlock block = cfgWorklist.remove(listSize);
				simulateBlock(block);
			}

			while (!cfgPhiWorklist.isEmpty()) {
				int listSize = cfgPhiWorklist.size() - 1;
				SsaBasicBlock block = cfgPhiWorklist.remove(listSize);
				simulatePhiBlock(block);
			}

			while (!varyingWorklist.isEmpty()) {
				int listSize = varyingWorklist.size() - 1;
				SsaInsn insn = varyingWorklist.remove(listSize);

				if (!executableBlocks.get(insn.getBlock().getIndex())) {
					continue;
				}

				if (insn instanceof PhiInsn) {
					simulatePhi((PhiInsn) insn);
				} else {
					simulateStmt(insn);
				}
			}
			while (!ssaWorklist.isEmpty()) {
				int listSize = ssaWorklist.size() - 1;
				SsaInsn insn = ssaWorklist.remove(listSize);

				if (!executableBlocks.get(insn.getBlock().getIndex())) {
					continue;
				}

				if (insn instanceof PhiInsn) {
					simulatePhi((PhiInsn) insn);
				} else {
					simulateStmt(insn);
				}
			}
		}

		replaceConstants();
		replaceBranches();
	}

	/**
	 * Replaces TypeBearers in source register specs with constant type bearers
	 * if possible. These are then referenced in later optimization steps.
	 */
	private void replaceConstants() {
		for (int reg = 0; reg < regCount; reg++) {
			if (latticeValues[reg] != CONSTANT) {
				continue;
			}
			if (!(latticeConstants[reg] instanceof TypedConstant)) {
				// We can't do much with these
				continue;
			}

			SsaInsn defn = ssaMeth.getDefinitionForRegister(reg);
			TypeBearer typeBearer = defn.getResult().getTypeBearer();

			if (typeBearer.isConstant()) {
				/*
				 * The definition was a constant already. The uses should be as
				 * well.
				 */
				continue;
			}

			// Update the destination RegisterSpec with the constant value
			RegisterSpec dest = defn.getResult();
			RegisterSpec newDest = dest
					.withType((TypedConstant) latticeConstants[reg]);
			defn.setResult(newDest);

			/*
			 * Update the sources RegisterSpec's of all non-move uses. These
			 * will be used in later steps.
			 */
			for (SsaInsn insn : ssaMeth.getUseListForRegister(reg)) {
				if (insn.isPhiOrMove()) {
					continue;
				}

				NormalSsaInsn nInsn = (NormalSsaInsn) insn;
				RegisterSpecList sources = insn.getSources();

				int index = sources.indexOfRegister(reg);

				RegisterSpec spec = sources.get(index);
				RegisterSpec newSpec = spec
						.withType((TypedConstant) latticeConstants[reg]);

				nInsn.changeOneSource(index, newSpec);
			}
		}
	}

	/**
	 * Replaces branches that have constant conditions with gotos
	 */
	private void replaceBranches() {
		for (SsaInsn insn : branchWorklist) {
			// Find if a successor block is never executed
			int oldSuccessor = -1;
			SsaBasicBlock block = insn.getBlock();
			int successorSize = block.getSuccessorList().size();
			for (int i = 0; i < successorSize; i++) {
				int successorBlock = block.getSuccessorList().get(i);
				if (!executableBlocks.get(successorBlock)) {
					oldSuccessor = successorBlock;
				}
			}

			/*
			 * Prune branches that have already been handled and ones that no
			 * longer have constant conditions (no nonexecutable successors)
			 */
			if (successorSize != 2 || oldSuccessor == -1)
				continue;

			// Replace branch with goto
			Insn originalRopInsn = insn.getOriginalRopInsn();
			block.replaceLastInsn(new PlainInsn(Rops.GOTO, originalRopInsn
					.getPosition(), null, RegisterSpecList.EMPTY));
			block.removeSuccessor(oldSuccessor);
		}
	}
}
